# @Desc    : do what 
# @Author  : meihongliang  
# @Time    : 2025-03-28 02:34
import os
import zipfile
import pandas as pd
from ftplib import FTP
import tempfile
from sqlalchemy import create_engine, Table, MetaData, insert
from sqlalchemy.exc import SQLAlchemyError
import numpy as np
import warnings

# 忽略特定的FutureWarning
warnings.simplefilter(action='ignore', category=FutureWarning)

# 配置信息
FTP_CONFIG = {
    'host': '192.168.1.4',
    'user': 'alpineftp_user',
    'passwd': 'alpineftp_password',
    'root_dir': '/ftp/alpineftp_user'
}

MYSQL_CONFIG = {
    'host': '192.168.1.4',
    'user': 'root',
    'password': 'root',
    'database': 'test_db',
    'table': 'test_kpi'
}

# 初始化SQLAlchemy引擎
engine = create_engine(
    f"mysql+pymysql://{MYSQL_CONFIG['user']}:{MYSQL_CONFIG['password']}@"
    f"{MYSQL_CONFIG['host']}/{MYSQL_CONFIG['database']}",
    pool_pre_ping=True,
    pool_recycle=3600,
    echo=False  # 关闭SQL日志输出
)


def connect_ftp():
    """安全连接FTP服务器"""
    try:
        ftp = FTP()
        ftp.connect(FTP_CONFIG['host'], 21, timeout=30)
        ftp.login(FTP_CONFIG['user'], FTP_CONFIG['passwd'])
        ftp.set_pasv(True)
        print(f"✓ FTP连接成功 (服务器: {FTP_CONFIG['host']})")
        return ftp
    except Exception as e:
        print(f"✗ FTP连接失败: {str(e)}")
        raise


def safe_insert_data(df, table_name):
    """
    使用SQLAlchemy核心API安全插入数据
    自动跳过重复主键记录（INSERT IGNORE）
    """
    if df.empty:
        print("! 空DataFrame，跳过插入")
        return

    # 确保列名合法
    df.columns = [str(col).strip().replace(' ', '_') for col in df.columns]

    metadata = MetaData()

    try:
        with engine.begin() as conn:  # 自动事务管理
            # 反射获取表结构
            table = Table(table_name, metadata, autoload_with=engine)

            # 分块处理大数据
            chunks = np.array_split(df, max(1, len(df) // 1000 + 1))
            total_inserted = 0

            for i, chunk in enumerate(chunks, 1):
                rows = chunk.to_dict('records')
                if rows:
                    stmt = insert(table).values(rows).prefix_with('IGNORE')
                    result = conn.execute(stmt)
                    inserted = result.rowcount
                    total_inserted += inserted
                    print(f"→ 批次 {i}: 尝试 {len(rows)} 行，实际插入 {inserted} 行")

            print(f"√ 完成！总共插入 {total_inserted}/{len(df)} 行数据")
    except SQLAlchemyError as e:
        print(f"✗ 数据库插入错误: {str(e)}")
        raise


def clean_dataframe(df):
    """数据清洗和预处理"""
    # 移除全空行
    df = df.dropna(how='all')

    # 确保没有转置操作警告
    if len(df.columns) > 100:  # 如果列太多可能误转置
        df = df.transpose() if len(df) < len(df.columns) else df

    # 类型转换示例
    for col in df.columns:
        if df[col].dtype == 'object':
            try:
                df[col] = pd.to_numeric(df[col], errors='ignore')
            except:
                pass

    return df


def process_zip_file(ftp, zip_path):
    """处理单个ZIP文件并返回解析出的DataFrame列表"""
    dfs = []

    with tempfile.NamedTemporaryFile(delete=False) as tmp_file:
        try:
            # 下载ZIP文件（带进度显示）
            print(f"\n↓ 下载ZIP: {zip_path}")

            def callback(data):
                tmp_file.write(data)

            ftp.retrbinary(f'RETR {zip_path}', callback)
            tmp_file.close()

            # 处理ZIP内容
            with zipfile.ZipFile(tmp_file.name, 'r') as zip_ref:
                for file in zip_ref.namelist():
                    # 忽略系统文件
                    if any(s in file for s in ['__MACOSX/', '.DS_Store', '._']):
                        continue

                    if file.lower().endswith('.csv'):
                        print(f"↳ 解析CSV: {file}")
                        try:
                            with zip_ref.open(file) as csv_file:
                                # 使用更安全的CSV读取方式
                                df = pd.read_csv(
                                    csv_file,
                                    encoding_errors='replace',
                                    on_bad_lines='warn'
                                )

                                df = clean_dataframe(df)

                                if not df.empty:
                                    dfs.append(df)
                                    print(f"✓ 解析成功: {len(df)} 行, {len(df.columns)} 列")
                        except Exception as e:
                            print(f"✗ 解析失败 [{file}]: {str(e)}")
                            with open('failed_files.log', 'a') as f:
                                f.write(f"{zip_path}/{file}\t{str(e)}\n")
        finally:
            os.unlink(tmp_file.name)

    return dfs


def scan_ftp_directory(ftp, path):
    """递归扫描FTP目录并处理ZIP文件"""
    try:
        print(f"\n📂 扫描目录: {path}")
        items = []
        ftp.cwd(path)
        ftp.dir(items.append)

        for item in items:
            if item.startswith('d'):
                # 处理子目录
                dirname = item.split()[-1]
                if dirname not in ('.', '..'):
                    scan_ftp_directory(ftp, f"{path}/{dirname}")
            else:
                # 处理文件
                filename = item.split()[-1]
                if filename.lower().endswith('.zip'):
                    full_path = f"{path}/{filename}"

                    try:
                        dfs = process_zip_file(ftp, full_path)
                        for df in dfs:
                            safe_insert_data(df, MYSQL_CONFIG['table'])
                    except Exception as e:
                        print(f"✗ 处理失败 [{full_path}]: {str(e)}")
    except Exception as e:
        print(f"✗ 目录扫描错误 [{path}]: {str(e)}")
        raise


def main():
    """主执行函数"""
    print("=" * 50)
    print("FTP到MySQL数据导入工具")
    print("=" * 50)

    ftp = None
    try:
        # 连接FTP
        ftp = connect_ftp()

        # 开始扫描处理
        scan_ftp_directory(ftp, FTP_CONFIG['root_dir'])

        print("\n" + "=" * 50)
        print("✅ 处理完成！")
        print("=" * 50)
    except KeyboardInterrupt:
        print("\n! 用户中断操作")
    except Exception as e:
        print("\n" + "=" * 50)
        print(f"❌ 处理中断: {str(e)}")
        print("=" * 50)
    finally:
        # 清理资源
        if ftp:
            ftp.quit()
        engine.dispose()
        print("所有连接已安全关闭")


if __name__ == '__main__':
    # 记录开始时间
    import time

    start_time = time.time()
    main()
    # 输出总耗时
    print(f"\n总耗时: {time.time() - start_time:.2f}秒")

    # TODO 根据文件夹名字往不同环境的mysql中存入数据
